package com.yth.utils.ql;

import com.ql.util.express.DefaultContext;
import com.ql.util.express.DynamicParamsUtil;
import com.ql.util.express.ExpressRunner;
import com.ql.util.express.InstructionSet;
import com.ql.util.express.config.QLExpressRunStrategy;
import com.ql.util.express.config.whitelist.CheckerFactory;
import com.ql.util.express.config.whitelist.WhiteChecker;
import com.ql.util.express.instruction.FunctionInstructionSet;
import com.ql.util.express.instruction.detail.*;
import com.ql.util.express.instruction.op.OperatorBase;
import com.ql.util.express.instruction.op.OperatorField;
import com.ql.util.express.instruction.op.OperatorMethod;
import com.yth.utils.ql.op.ArrayStreamOperatorFactory;
import com.yth.utils.ql.op.NewLinkedMapOperator;
import com.yth.utils.udf.Udf;
import com.yth.utils.udf.UdfLoader;
import lombok.EqualsAndHashCode;
import lombok.Getter;

import java.lang.reflect.Field;
import java.text.DateFormat;
import java.text.NumberFormat;
import java.time.*;
import java.time.format.DateTimeFormatter;
import java.util.*;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.BaseStream;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * QL表达式帮助类
 *
 * @author yutianhong
 * @version 1.0
 * @since 2024/12/13 14:36
 */
public class QLHelper {

    private static final ExpressRunner EXPRESS_RUNNER = initializeExpressRunner();
    private static final int TIMEOUT_MILLIS = 1500;
    private static final String IMPORT_STRING = buildUDFImportString();
    private static final Field LAMBDA_SET_FIELD = initializeLambdaSetField();

    static {
        setupQLExpressRunStrategy();
    }

    /**
     * Executes the provided expression with the given parameters.
     */
    public static Object execute(String expression, Map<String, ?> param) throws Exception {
        if (expression == null) {
            throw new IllegalArgumentException("expression is null");
        }
        DefaultContext<String, Object> context = new DefaultContext<>();
        if (param != null) {
            context.putAll(param);
        }
        return EXPRESS_RUNNER.execute(IMPORT_STRING + expression, context, null, true, false, TIMEOUT_MILLIS);
    }

    public static void tryCompile(String text) throws Exception {
        EXPRESS_RUNNER.parseInstructionSet(IMPORT_STRING + text);
    }

    /**
     * 这个方法不能保证能够拿到所有参数。
     */
    public static String[] getOutVarNames(String expression) throws Exception {
        return EXPRESS_RUNNER.getOutVarNames(IMPORT_STRING + expression);
    }

    /**
     * 提取表达式中的两级外部变量。
     * <p>注意：会包含局部变量。
     * <p>Examples of usage:
     * <ul>
     *   <li>{@code getVariableNames("a")} returns an empty set: {@code []}</li>
     *   <li>{@code getVariableNames("ds.save")} returns: {@code [[ds, save]]}</li>
     *   <li>{@code getVariableNames("ds.save + ds.list")} returns: {@code [[ds, save], [ds, list]]}</li>
     *   <li>{@code getVariableNames("ds.get("save")")} returns: {@code [[ds, save]]}</li>
     *   <li>{@code getVariableNames("a.stream().map(e -> e.field = ds.save).collect(Collectors.toList())")}
     *       returns: {@code [[ds, save], [e, field]]}</li>
     *  <li>{@code getVariableNames("function add(int i) {return i + ds.save} return add(1);")}
     *       returns: {@code [[ds, save]]}</li>
     * </ul>
     *
     * @param expression 表达式
     * @return 变量集合
     * @throws Exception 表达式语法错误
     */
    public static Set<VariablePair> getVariablePairs(String expression) throws Exception {
        InstructionSet instructionSet = EXPRESS_RUNNER.parseInstructionSet(IMPORT_STRING + expression);
        Set<VariablePair> allVars = new HashSet<>();
        travelInstructionSet(instructionSet, allVars);
        return allVars;
    }

    @Getter
    @EqualsAndHashCode
    public static class VariablePair {
        private final String varName;
        private final String attributeName;

        public VariablePair(String varName, String attributeName) {
            this.varName = varName;
            this.attributeName = attributeName;
        }

        @Override
        public String toString() {
            return "[" + varName + ", " + attributeName + ']';
        }
    }

    private static void travelInstructionSet(InstructionSet instructionSet, Set<VariablePair> allVars) throws Exception {
        if (instructionSet == null) {
            return; // paranoid
        }
        FunctionInstructionSet[] functions = instructionSet.getFunctionInstructionSets();
        if (functions != null) {
            for (FunctionInstructionSet func : functions) {
                // function body
                travelInstructionSet(func.instructionSet, allVars);  // recurse
            }
        }
        int instructionLength = instructionSet.getInstructionLength();
        for (int i = 0; i < instructionLength; i++) {
            Instruction instruction = instructionSet.getInstruction(i);
            if (instruction instanceof InstructionLoadLambda) {
                // lambda expression
                InstructionLoadLambda instructionLoadLambda = (InstructionLoadLambda) instruction;
                Object object = LAMBDA_SET_FIELD.get(instructionLoadLambda);
                if (object instanceof InstructionSet) {
                    InstructionSet lambdaInstructionSet = (InstructionSet) object;
                    travelInstructionSet(lambdaInstructionSet, allVars);  // recurse
                }
            } else if (instruction instanceof InstructionLoadAttr) {
                InstructionLoadAttr instructionLoadAttr = (InstructionLoadAttr) instruction;
                String attr = instructionLoadAttr.getAttrName();
                if ("null".equals(attr) || i + 1 >= instructionLength) {
                    continue;
                }
                Instruction next = instructionSet.getInstruction(i + 1);
                if (next instanceof InstructionOperator) {
                    // situation1: var.field
                    OperatorBase opBase = ((InstructionOperator) next).getOperator();
                    if (opBase instanceof OperatorField) {
                        // tick to get filedName
                        String fieldName = opBase.toString().split(":")[1];
                        VariablePair vars = new VariablePair(attr, fieldName);
                        allVars.add(vars);
                    }
                } else if (next instanceof InstructionConstData && i + 2 < instructionLength) {
                    // situation2: var.get("field")
                    InstructionConstData constData = (InstructionConstData) next;
                    Object objectInner = constData.getOperateData().getObjectInner(null);
                    Instruction afterNext = instructionSet.getInstruction(i + 2);
                    if (objectInner instanceof String && afterNext instanceof InstructionOperator) {
                        OperatorBase opBase = ((InstructionOperator) afterNext).getOperator();
                        if (opBase instanceof OperatorMethod) {
                            String methodName = opBase.toString().split(":")[1];
                            if ("get".equals(methodName)) {
                                VariablePair vars = new VariablePair(attr, (String) objectInner);
                                allVars.add(vars);
                            }
                        }
                    }
                }
            }
        }
    }

    private static void setupQLExpressRunStrategy() {
        // 运行期黑名单设置
        QLExpressRunStrategy.setForbidInvokeSecurityRiskMethods(true);
        QLExpressRunStrategy.setForbidInvokeSecurityRiskConstructors(true);
        QLExpressRunStrategy.addSecurityRiskMethod(System.class, "gc");
        QLExpressRunStrategy.addSecurityRiskMethod(BaseStream.class, "parallel");
        // 编译期白名单设置
        List<WhiteChecker> checkersOfJava = buildWhiteCheckers();
        checkersOfJava.add(CheckerFactory.assignable(Udf.class));   // support udf
        QLExpressRunStrategy.setCompileWhiteCheckerList(checkersOfJava);
        // 限制最大循环次数
//        QLExpressRunStrategy.setMaxArrLength(100);
        // 支持函数动态参数调用
        DynamicParamsUtil.supportDynamicParams = true;
    }

    private static ExpressRunner initializeExpressRunner() {
        ExpressRunner runner = new ExpressRunner(true, false);
        runner.setShortCircuit(true);
        // add `stream()` method for array class - 使用专门的操作符实例
        runner.addClassMethod("stream", Object[].class, ArrayStreamOperatorFactory.createObjectArrayStreamOperator());
        runner.addClassMethod("stream", int[].class, ArrayStreamOperatorFactory.createIntArrayStreamOperator());
        runner.addClassMethod("stream", long[].class, ArrayStreamOperatorFactory.createLongArrayStreamOperator());
        runner.addClassMethod("stream", double[].class, ArrayStreamOperatorFactory.createDoubleArrayStreamOperator());
        runner.addClassMethod("stream", byte[].class, ArrayStreamOperatorFactory.createByteArrayStreamOperator());
        runner.addClassMethod("stream", boolean[].class, ArrayStreamOperatorFactory.createBooleanArrayStreamOperator());
        runner.addClassMethod("stream", short[].class, ArrayStreamOperatorFactory.createShortArrayStreamOperator());
        runner.addClassMethod("stream", char[].class, ArrayStreamOperatorFactory.createCharArrayStreamOperator());
        runner.addClassMethod("stream", float[].class, ArrayStreamOperatorFactory.createFloatArrayStreamOperator());
        // add `NewLinkedMap()` function
        runner.addFunction("NewLinkedMap", new NewLinkedMapOperator("NewLinkedMap"));
        return runner;
    }

    private static Field initializeLambdaSetField() {
        try {
            Field field = InstructionLoadLambda.class.getDeclaredField("lambdaSet");
            field.setAccessible(true);
            return field;
        } catch (Exception e) {
            throw new RuntimeException("Failed to initialize lambda set field", e);
        }
    }

    private static List<WhiteChecker> buildWhiteCheckers() {
        List<WhiteChecker> checkersOfJava = new ArrayList<>(Arrays.asList(
                // 字符串
                CheckerFactory.assignable(CharSequence.class),
                CheckerFactory.must(StringTokenizer.class),
                CheckerFactory.must(StringJoiner.class),
                CheckerFactory.must(Pattern.class),
                CheckerFactory.must(Matcher.class),
                // 基础类型
                CheckerFactory.must(Object.class),
                CheckerFactory.must(Void.class),
                CheckerFactory.must(Boolean.class),
                CheckerFactory.must(Character.class),
                CheckerFactory.assignable(Number.class),
                CheckerFactory.assignable(NumberFormat.class),
                // 日期处理
                CheckerFactory.assignable(Date.class),
                CheckerFactory.must(Calendar.class),
                CheckerFactory.assignable(DateFormat.class),
                CheckerFactory.assignable(DateTimeFormatter.class),
                CheckerFactory.must(Duration.class),
                CheckerFactory.must(LocalDateTime.class),
                CheckerFactory.must(ZonedDateTime.class),
                CheckerFactory.must(Instant.class),
                CheckerFactory.must(Month.class),
                CheckerFactory.must(Year.class),
                CheckerFactory.must(LocalDate.class),
                CheckerFactory.must(LocalTime.class),
                CheckerFactory.must(ZoneId.class),
                CheckerFactory.must(MonthDay.class),
                CheckerFactory.must(YearMonth.class),

                // 集合
                CheckerFactory.assignable(List.class),
                CheckerFactory.assignable(Set.class),
                CheckerFactory.assignable(Map.class),
                CheckerFactory.must(Iterator.class),
                CheckerFactory.must(BitSet.class),
                CheckerFactory.must(Arrays.class),
                CheckerFactory.must(Collections.class),
                CheckerFactory.must(Collectors.class),
                CheckerFactory.must(Stream.class),
                // 其他
                CheckerFactory.must(Math.class),
//                CheckerFactory.must(Thread.class),
                CheckerFactory.must(System.class),
                CheckerFactory.must(Objects.class),
                CheckerFactory.must(Optional.class),
                CheckerFactory.must(OptionalInt.class),
                CheckerFactory.must(OptionalLong.class),
                CheckerFactory.must(OptionalDouble.class),
                CheckerFactory.must(Base64.class),
                CheckerFactory.must(Enum.class),
                CheckerFactory.assignable(Throwable.class),
                CheckerFactory.assignable(Error.class),
                CheckerFactory.assignable(StackTraceElement.class)
        ));

//        try {
//            checkersOfJava.addAll(Arrays.asList(
//                    CheckerFactory.assignable(com.alibaba.fastjson.JSON.class),
//                    CheckerFactory.must(com.alibaba.fastjson.JSONPath.class)));
//        } catch (Throwable ignore) {
//        }
//        try {
//            checkersOfJava.add(CheckerFactory.assignable(ObjectMapper.class));
//        } catch (Throwable ignore) {
//        }
        // 三方工具
//        try {
//            checkersOfJava.add(CheckerFactory.assignable(DigestUtils.class));
//        } catch (Throwable ignore) {
//        }
//        try {
//            checkersOfJava.addAll(Arrays.asList(
//                    CheckerFactory.assignable(org.apache.commons.lang3.ArrayUtils.class),
//                    CheckerFactory.assignable(org.apache.commons.lang3.RegExUtils.class),
//                    CheckerFactory.assignable(org.apache.commons.lang3.BooleanUtils.class),
//                    CheckerFactory.assignable(org.apache.commons.lang3.StringUtils.class),
//                    CheckerFactory.assignable(org.apache.commons.lang3.math.NumberUtils.class)));
//        } catch (Throwable ignore) {
//        }
        return checkersOfJava;
    }

    private static String buildUDFImportString() {
        List<Class<? extends Udf>> udfClasses = UdfLoader.getUdfClasses();
        StringBuilder stringBuilder = new StringBuilder();
        for (Class<? extends Udf> udfClass : udfClasses) {
            stringBuilder.append("import ")
                    .append(udfClass.getName())
                    .append(";\n");
        }
        return stringBuilder.toString();
    }

    private QLHelper() {
        // Prevent instantiation
    }
}
